package com.enba.boot.core.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;

public class TreeUtils {

  /**
   * 构建树形结构数据的方法
   * 该方法接收两个参数：一个用于获取扁平化数据列表的供应商函数，另一个用于将扁平化数据转换为树形结构节点的函数
   * 通过这些参数，方法能够灵活地从不同的数据源获取数据，并将其转换为所需的树形结构
   *
   * @param supplierFlattenData 一个Supplier函数，用于提供扁平化数据列表
   * @param convertTreeData 一个Function函数，用于将扁平化数据列表转换为树形结构节点列表
   * @return 返回构建好的树形结构数据列表
   */
  public static <P> List<TreeNode> buildTree(
      Supplier<List<P>> supplierFlattenData, Function<List<P>, List<TreeNode>> convertTreeData) {

    // 1.业务方提供扁平化节点数据，建议直接从数据库查询
    List<P> flattenDataList = supplierFlattenData.get();

    // 2.业务方提供处理扁平化节点数据方法，进一步处理步骤1获得的数据
    List<TreeNode> applyList = convertTreeData.apply(flattenDataList);

    // 3.返回前端需要的树形结构数据
    return TreeUtils.buildTree(applyList);
  }

  /**
   * 将父子结构的数据转换为树形结构
   *
   * @param nodes 所有节点数据
   * @return 树形结构的根节点列表
   */
  public static List<TreeNode> buildTree(List<? extends TreeNode> nodes) {
    // 用于存储所有节点的映射表（key: id, value: TreeNode）
    Map<Long, TreeNode> nodeMap = new HashMap<>();

    // 用于存储根节点
    List<TreeNode> rootNodes = new ArrayList<>();

    // 将所有节点放入映射表
    for (TreeNode node : nodes) {
      nodeMap.put(node.getId(), node);
    }

    // 构建树结构
    for (TreeNode node : nodes) {
      Long parentId = node.getParentId();
      if (parentId == null || parentId == 0) {
        // 如果 parentId 为 null 或 0，说明是根节点
        rootNodes.add(node);
      } else {
        // 否则，找到父节点，并将当前节点添加到父节点的 children 中
        TreeNode parentNode = nodeMap.get(parentId);
        if (parentNode != null) {
          parentNode.getChildren().add(node);
        }
      }
    }

    return rootNodes;
  }

  public static void main(String[] args) {
    // 示例数据
    List<TreeNode> nodes = new ArrayList<>();
    nodes.add(new TreeNode(1L, null, "总公司"));
    nodes.add(new TreeNode(2L, 1L, "分公司A"));
    nodes.add(new TreeNode(3L, 1L, "分公司B"));
    nodes.add(new TreeNode(4L, 2L, "部门A1"));
    nodes.add(new TreeNode(5L, 2L, "部门A2"));
    nodes.add(new TreeNode(6L, 3L, "部门B1"));

    // 生成树结构
    List<TreeNode> tree = buildTree(nodes);

    // 打印树结构
    for (TreeNode root : tree) {
      printTree(root, 0);
    }
  }

  /**
   * 打印树结构（递归）
   *
   * @param node 当前节点
   * @param level 当前层级（用于缩进）
   */
  private static void printTree(TreeNode node, int level) {
    for (int i = 0; i < level; i++) {
      System.out.print("--");
    }
    System.out.println(node.getName());

    for (TreeNode child : node.getChildren()) {
      printTree(child, level + 1);
    }
  }
}
